#!/usr/bin/env python
# coding: utf-8

# In[50]:


import torch
from torch import nn

from labml.logger import inspect
from py.attention_impl import MultiHeadAttention



# 传统Transformer模型使用绝对位置编码，即每个位置有一个固定的编码向量。然而，这种方法在处理长序列时可能无法有效捕捉相对位置关系。ROPE通过引入旋转矩阵，将位置信息嵌入到模型的注意力机制中，从而更好地建模相对位置关系。
# 
# ROPE的核心思想是通过旋转矩阵对词向量进行变换，将位置信息融入词向量中。假设输入向量的长度是偶数（一般也确实是偶数）：
# 
# $$
# \begin{align}
# RoPE(\bf{x}_m,m) &=\bf{R_m}\bf{x}_m\\
# &=\begin{pmatrix}
# \cos m\theta_1 & -\sin m\theta_1 & 0 &\cdots & \cdots & 0 \\
# \sin m\theta_1 & \cos m\theta_1 & 0 &\cdots & \cdots & 0 \\
# 0 & 0 & \cos m\theta_2 & -\sin m\theta_2 & \cdots & 0\\
# 0 & 0 & \sin m\theta_2 & \cos m\theta_2& \cdots & 0 \\
# \vdots & \vdots & \cdots &\ddots & \cos m\theta_{\frac{d}{2}} & -\sin m\theta_{\frac{d}{2}}\\
# 0 & 0 & \cdots &\cdots & \sin m\theta_{\frac{d}{2}} & \cos m\theta_{\frac{d}{2}}
# \end{pmatrix}
# \begin{pmatrix}
# x_m^{(1)} \\
# x_m^{(1+\frac{d}{2})} \\
# \vdots\\
# x_m^{(\frac{d}{2})}\\
# x_m^{(d)}
# \end{pmatrix}
# \end{align}
# 
# $$
# 式中，$\theta_i=10000^{\frac{-2i}{n}}$，m是输入向量对应token的从0开始的索引位置。完成计算后，各个维度的位置还应当还原，因此在计算上，可以直接在原值上计算。
# 
# 对于第i对（从1开始数），对应的输出结果是：
# $$
# \begin{pmatrix}
# x_m^{(i)}\cos m\theta -x_m^{(i+\frac{d}{2})}\sin m\theta \\
# x_m^{(i)}\sin m\theta + x_m^{(i+\frac{d}{2})}\cos m\theta
# \end{pmatrix}
# $$
# 注意到，RoPE并不是计算相邻2个维度的旋转结果，而是隔d/2组成一对，这可以降低计算复杂度，增加维度分布的平均性，防止某个维度的信息被忽略
# 
# 可以看到，RoPE的核心思想是把输入向量中每2个分量视为1个二维平面，并分别应用旋转矩阵。这样设计的核心原因在于，可以将两个不同绝对位置对应分量的点积，转换成相对位置的计算：
# $$
# \begin{align}
# &\langle RoPE(\vec{x}_m, m), RoPE(\vec{x}_n , n) \rangle\\
# &=\langle RoPE(\vec{x}_m, m-n), RoPE(\vec{x}_n , 0) \rangle
# \end{align}
# $$
# 

# 传统Transformer模型使用绝对位置编码，即每个位置有一个固定的编码向量。然而，这种方法在处理长序列时可能无法有效捕捉相对位置关系。ROPE通过引入旋转矩阵，将位置信息嵌入到模型的注意力机制中，从而更好地建模相对位置关系。
# 
# ROPE的核心思想是通过旋转矩阵对词向量进行变换，将位置信息融入词向量中。假设输入向量的长度是偶数（一般也确实是偶数）：
# 
# $$
# \begin{align}
# RoPE(\bf{x}_m,m) &=\bf{R_m}\bf{x}_m\\
# &=\begin{pmatrix}
# \cos m\theta_1 & -\sin m\theta_1 & 0 &\cdots & \cdots & 0 \\
# \sin m\theta_1 & \cos m\theta_1 & 0 &\cdots & \cdots & 0 \\
# 0 & 0 & \cos m\theta_2 & -\sin m\theta_2 & \cdots & 0\\
# 0 & 0 & \sin m\theta_2 & \cos m\theta_2& \cdots & 0 \\
# \vdots & \vdots & \cdots &\ddots & \cos m\theta_{\frac{d}{2}} & -\sin m\theta_{\frac{d}{2}}\\
# 0 & 0 & \cdots &\cdots & \sin m\theta_{\frac{d}{2}} & \cos m\theta_{\frac{d}{2}}
# \end{pmatrix}
# \begin{pmatrix}
# x_m^{(1)} \\
# x_m^{(1+\frac{d}{2})} \\
# \vdots\\
# x_m^{(\frac{d}{2})}\\
# x_m^{(d)}
# \end{pmatrix}
# \end{align}
# 
# $$
# 式中，$\theta_i=10000^{\frac{-2i}{n}}$，m是输入向量对应token的从0开始的索引位置。
# 
# 对于第i对（从1开始数），对应的输出结果是：
# $$
# \begin{pmatrix}
# x_m^{(i)}\cos m\theta -x_m^{(i+\frac{d}{2})}\sin m\theta \\
# x_m^{(i)}\sin m\theta + x_m^{(i+\frac{d}{2})}\cos m\theta
# \end{pmatrix}
# $$
# 注意到，RoPE并不是计算相邻2个维度的旋转结果，而是隔d/2组成一对，这可以降低计算复杂度，增加维度分布的平均性，防止某个维度的信息被忽略
# 
# 可以看到，RoPE的核心思想是把输入向量中每2个分量视为1个二维平面，并分别应用旋转矩阵。这样设计的核心原因在于，可以将两个不同绝对位置对应分量的点积，转换成相对位置的计算：
# $$
# \begin{align}
# &\langle RoPE(\vec{x}_m, m), RoPE(\vec{x}_n , n) \rangle\\
# &=\langle RoPE(\vec{x}_m, m-n), RoPE(\vec{x}_n , 0) \rangle
# \end{align}
# $$
# 

# In[ ]:


class RotaryPositionalEmbeddings(nn.Module):
    def __init__(self, d: int, base: int = 10_000):
        super().__init__()
        self.base = base
        self.d = d
        self._cos_cached = self._sin_cached = None

    def _build_cached(self, x: torch.Tensor):
        # 已经build的cache的大小大于已有dimension的大小
        if self._cos_cached is not None and x.shape[0] <= self._cos_cached.shape[0]:
            return
        seq_len = x.shape[0]
        # 见上述公式
        theta = (1.0 / self.base **
                 (torch.arange(0, self.d, 2).float()/self.d)).to(x.device)


        seq_idx = torch.arange(seq_len, device=x.device, dtype=torch.float32)
        # 长度为n，维度为d/2的数据，就需要nd/2个角度值,shape = (seq_len, d//2)
        # 这样表示的einsum就是在计算外积
        # 实际的数值中，idx_angle[i][j] = i*θ_j
        idx_angle = torch.einsum("n,d->nd", seq_idx, theta)
        # shape = (seq_len, d)
        both_idx_angle = torch.cat([idx_angle, idx_angle], dim=1)
        # shape = (seq_len, 1,1, d)，用于将计算结果广播
        self._sin_cached = both_idx_angle.sin()[:, None, None, :]
        self._cos_cached = both_idx_angle.cos()[:, None, None, :]


    def _neg_half(self, x: torch.Tensor) -> torch.Tensor:
        """将需要rope的x部分按最后的dmodel维度，前半段设为-,后半段保持原样,并将前后半段旋转

        Args:
            x (torch.Tensor): _description_
        """
        d_2 = self.d // 2
        return torch.cat([-x[..., d_2:], x[..., :d_2]], dim=-1)

    def forward(self, x: torch.Tensor):
        self._build_cached(x)
        seq_len = x.shape[0]
        assert x.shape[-1] % self.d == 0
        # RoPE的另一特性，分组旋转
        # chatGLM则选择只旋转一部分，剩下的原样输出，减少计算量
        x_chunk = torch.chunk(x, x.shape[-1] // self.d, dim=-1)
        x_rope_chunk = []
        for x_rope in x_chunk:
            neg_half_x = self._neg_half(x_rope)

            # 根据旋转公式计算, 这里不是直接计算出矩阵行*向量列直接得到一个结果分量，
            # 而是用矩阵列分别乘以向量分量，分别计算每一个矩阵列对最终运算的结果后再加和
            x_rope_chunk.append( x_rope * self._cos_cached[:seq_len] + neg_half_x * self._sin_cached[:seq_len])
        return torch.cat(x_rope_chunk, dim=-1)


# In[ ]:


from tkinter import NO
import torch


class MyRoPE(nn.Module):
    def __init__(self, d: int, base:int = 10000):
        super().__init__()
        self.d = d
        self.base = base
        self._sin_cached = self._cos_cached = None


    def _build_cache(self, x: torch.Tensor):
        build_len = (x.shape[0]+1)//2*2
        if self._sin_cached is not None and self._sin_cached.shape[0] >= build_len:
            return

        exponent = -torch.arange(0,self.d,2, device=x.device).float()/self.d
        theta = torch.pow(self.base, exponent)
        seq_idx = torch.arange(build_len, dtype=x.dtype, device=x.device)

        angle = torch.einsum("n,d->nd", seq_idx,theta)
        angle = torch.cat([angle,angle], dim=-1)

        self._sin_cached = angle.sin()[:,None,None,:]
        self._cos_cached = angle.cos()[:,None,None,:]

    def forward(self, x: torch.Tensor, inplace=True):
        self._build_cache(x)
        seq_len,*_, model_dim = x.shape
        chunked_x = x.chunk(model_dim//self.d, -1)
        chunks = [] 
        for id,chunk in enumerate(chunked_x):
            neg_half = torch.cat([-chunk[...,self.d//2:], chunk[...,:self.d//2]],dim=-1)
            mid_val = chunk*self._cos_cached[:seq_len] + neg_half*self._sin_cached[:seq_len]
            if inplace:
                x[...,id*self.d:(id+1)*self.d] = mid_val
            else:
                chunks.append(mid_val)
        return x if inplace else torch.cat(chunks, dim=-1)


heads, d_model, seq_len_q, batch = 1, 4, 4, 1
my_rope = MyRoPE(d_model//heads)
x = torch.randn((seq_len_q, batch,heads, d_model//heads))

rope = RotaryPositionalEmbeddings(d_model//heads)
# print(rope(x).shape)
# print(my_rope(x).shape)
print(torch.allclose(rope(x), my_rope(x,False)))









# In[28]:


class RotaryPEMultiHeadAttention(MultiHeadAttention):
    def __init__(self, heads: int, d_model: int, rope_chunk: int = 2, dropout_prob: float = 0.0):
        super().__init__(heads, d_model, dropout_prob)
        self.query_rope = RotaryPositionalEmbeddings(d_model//rope_chunk)
        self.key_rope = RotaryPositionalEmbeddings(d_model//rope_chunk)

    def get_score(self, query: torch.Tensor, key: torch.Tensor):
        return torch.einsum("qbhd,kbhd->qkbh", self.query_rope(query), self.key_rope(key))


# In[35]:


def _test_rotary():
    x = torch.tensor([[1, 2, 3, 4], [4, 5, 6, 7], [
                     7, 8, 9, 10]], dtype=torch.float)
    x = x[:, None, None, :]
    print(x)

    rotary_pe = RotaryPositionalEmbeddings(4)
    print(rotary_pe(x))


if __name__ == '__main__':
    _test_rotary()

